"""
加密货币数据采集器模块

负责通过akshare API获取加密货币数据，进行清洗，并保存到数据库
"""
import time
import datetime
import logging
import akshare as ak
import pandas as pd
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from config.settings import DATA_COLLECTION

# 配置日志
logger = logging.getLogger("crypto_collector")


class CryptoCollector:
    """加密货币数据采集类"""
    
    def __init__(self):
        """初始化采集器"""
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.batch_size = DATA_COLLECTION["batch_size"]
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def get_crypto_list(self, limit=200):
        """
        获取加密货币列表
        
        Args:
            limit: 返回的加密货币数量，默认200
            
        Returns:
            加密货币列表DataFrame
        """
        try:
            logger.info(f"开始获取加密货币列表，数量限制: {limit}")
            
            # 使用akshare获取加密货币行情
            crypto_df = ak.crypto_hist(symbol="usdt")
            
            # 获取市值排名靠前的加密货币
            crypto_df = crypto_df.sort_values(by="market_cap", ascending=False).head(limit)
            
            # 提取需要的列
            result_df = pd.DataFrame({
                "code": crypto_df["symbol"],
                "name": crypto_df["name"],
                "market_cap": crypto_df["market_cap"],
                "price": crypto_df["price"]
            })
            
            logger.info(f"成功获取 {len(result_df)} 种加密货币信息")
            return result_df
            
        except Exception as e:
            logger.error(f"获取加密货币列表失败: {str(e)}")
            raise
    
    def collect_crypto_info(self, limit=200):
        """
        采集并更新加密货币基本信息
        
        Args:
            limit: 采集数量限制，默认200
            
        Returns:
            (新增数量, 更新数量)
        """
        logger.info(f"开始采集加密货币基本信息")
        
        try:
            # 获取加密货币列表
            crypto_df = self.get_crypto_list(limit)
            
            # 获取已存在的加密货币代码
            existing_codes = {
                code[0] for code in 
                self.db.query(SecurityInfo.code).filter(SecurityInfo.type == 'crypto').all()
            }
            
            # 准备批量插入或更新的数据
            new_cryptos = []
            update_cryptos = []
            
            for _, row in crypto_df.iterrows():
                code = row['code'].upper()  # 统一使用大写代码
                name = row['name']
                
                # 检查是否存在
                if code in existing_codes:
                    # 更新现有记录
                    crypto = self.db.query(SecurityInfo).filter(
                        SecurityInfo.code == code,
                        SecurityInfo.type == 'crypto'
                    ).first()
                    
                    if crypto:
                        crypto.name = name
                        crypto.update_time = datetime.datetime.now()
                        update_cryptos.append(crypto)
                else:
                    # 创建新记录
                    new_crypto = SecurityInfo(
                        code=code,
                        name=name,
                        type='crypto',
                        exchange='CRYPTO',
                        is_active=1
                    )
                    new_cryptos.append(new_crypto)
            
            # 批量保存新记录
            if new_cryptos:
                self.db.add_all(new_cryptos)
                self.db.commit()
                logger.info(f"成功添加 {len(new_cryptos)} 种新加密货币")
            
            # 批量更新现有记录
            if update_cryptos:
                self.db.commit()
                logger.info(f"成功更新 {len(update_cryptos)} 种加密货币信息")
            
            return len(new_cryptos), len(update_cryptos)
            
        except Exception as e:
            self.db.rollback()
            logger.error(f"采集加密货币基本信息失败: {str(e)}")
            raise
    
    def collect_crypto_history(self, symbol, start_date=None, end_date=None, task_id=None):
        """
        采集指定加密货币的历史数据
        
        Args:
            symbol: 加密货币代码
            start_date: 开始日期，默认为None
            end_date: 结束日期，默认为None
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            self._update_task_status(task_id, "running")
        
        # 获取加密货币信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == symbol.upper(),
            SecurityInfo.type == 'crypto'
        ).first()
        
        if not security:
            error_msg = f"加密货币 {symbol} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
        
        # 设置日期范围
        if not start_date:
            # 从最后一条记录的下一天开始
            last_record = self.db.query(DailyQuote).filter(
                DailyQuote.security_id == security.id
            ).order_by(DailyQuote.trade_date.desc()).first()
            
            if last_record:
                start_date = (last_record.trade_date + datetime.timedelta(days=1)).strftime('%Y%m%d')
            else:
                # 默认采集近1年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365)).strftime('%Y%m%d')
        
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        # 转换日期格式
        if isinstance(start_date, str):
            start_date = pd.to_datetime(start_date).strftime('%Y-%m-%d')
        if isinstance(end_date, str):
            end_date = pd.to_datetime(end_date).strftime('%Y-%m-%d')
        
        logger.info(f"开始采集加密货币 {symbol} 的历史数据，时间范围: {start_date} 至 {end_date}")
        
        try:
            # 重试机制
            retry_count = 0
            df = None
            
            while retry_count < self.retry_limit:
                try:
                    # 使用akshare获取历史数据
                    df = ak.crypto_hist(symbol=symbol.lower(), period="daily", start_date=start_date, end_date=end_date)
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning(f"采集加密货币 {symbol} 数据失败，重试 {retry_count}/{self.retry_limit}: {str(e)}")
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df is None or df.empty:
                logger.info(f"加密货币 {symbol} 在指定时间范围内没有数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
            
            # 清洗和转换数据
            df = self._clean_crypto_data(df)
            
            # 保存到数据库
            inserted_count = self._save_crypto_quotes(security.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集加密货币 {symbol} 数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
    
    def _clean_crypto_data(self, df):
        """清洗和标准化加密货币数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.lower() for col in df.columns]
        
        # 加密货币数据的列名映射
        name_mapping = {
            'date': 'trade_date',
            'symbol': 'code',
            'open': 'open',
            'high': 'high',
            'low': 'low',
            'close': 'close',
            'volume': 'volume',
            'market_cap': 'market_cap'
        }
        
        # 应用列名映射
        df = df.rename(columns={k: v for k, v in name_mapping.items() if k in df.columns})
        
        # 处理日期
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date'])
        
        # 处理数值类型
        numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'market_cap']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 保留需要的列
        keep_cols = [col for col in ['trade_date', 'open', 'high', 'low', 'close', 
                                     'volume', 'market_cap'] if col in df.columns]
        
        return df[keep_cols]
    
    def _save_crypto_quotes(self, security_id, df):
        """保存加密货币行情数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的交易日期
            existing_dates = {
                date[0] for date in 
                self.db.query(DailyQuote.trade_date).filter(
                    DailyQuote.security_id == security_id
                ).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                trade_date = row['trade_date'].date()
                
                # 跳过已存在的记录
                if trade_date in existing_dates:
                    continue
                
                daily_quote = DailyQuote(
                    security_id=security_id,
                    trade_date=trade_date,
                    open=row.get('open'),
                    high=row.get('high'),
                    low=row.get('low'),
                    close=row.get('close'),
                    adj_close=row.get('close'),  # 加密货币没有复权价
                    volume=row.get('volume'),
                    amount=None,  # 加密货币可能没有成交额
                    turnover=None  # 加密货币没有换手率
                )
                
                batch_records.append(daily_quote)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info(f"成功保存 {inserted_count} 条加密货币日线数据")
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存加密货币日线数据失败(完整性错误): {str(e)}")
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存加密货币日线数据失败: {str(e)}")
            raise
    
    def _update_task_status(self, task_id, status, error_message=None, result_count=None):
        """更新任务状态"""
        try:
            task = self.db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
            if task:
                task.status = status
                if error_message:
                    task.error_message = error_message
                if result_count is not None:
                    task.result_count = result_count
                task.update_time = datetime.datetime.now()
                self.db.commit()
        except Exception as e:
            self.db.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")


# 测试函数
def test_crypto_collector():
    """测试加密货币数据采集器功能"""
    collector = CryptoCollector()
    
    # 1. 测试获取加密货币列表
    try:
        crypto_list = collector.get_crypto_list(limit=20)
        print(f"成功获取 {len(crypto_list)} 种加密货币信息")
        print(crypto_list.head())
    except Exception as e:
        print(f"获取加密货币列表失败: {str(e)}")
    
    # 2. 测试采集加密货币基本信息
    try:
        new_count, update_count = collector.collect_crypto_info(limit=20)
        print(f"新增加密货币: {new_count}, 更新加密货币: {update_count}")
    except Exception as e:
        print(f"采集加密货币基本信息失败: {str(e)}")
    
    # 3. 测试采集历史数据
    if crypto_list is not None and not crypto_list.empty:
        # 选择第一个加密货币进行测试
        test_symbol = crypto_list.iloc[0]['code']
        try:
            # 采集最近30天数据
            end_date = datetime.datetime.now().strftime('%Y-%m-%d')
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y-%m-%d')
            
            count = collector.collect_crypto_history(test_symbol, start_date, end_date)
            print(f"加密货币 {test_symbol} 成功采集 {count} 条历史数据")
        except Exception as e:
            print(f"采集加密货币 {test_symbol} 历史数据失败: {str(e)}")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_crypto_collector()